#!/usr/bin/env python

import os, sys
from astropy.io import fits
import numpy as np
from argparse import ArgumentParser


# names and coordiantes of the irradiated regions
reg_names = ['A (DD)', 'B (TID)', 'C (TID)', 'D (DD)', 'E (NONE)', 'F (NONE)']
reg = [ [206,  4506, 182,  4382],       # A
        [4698, 8998, 168,  4368],       # B
        [4993, 7767, 5009, 9059],       # C
        [1444, 4344, 4859, 9073],       # D
        [8222, 9078, 4903, 9703],       # E
        [82,   1172, 4717, 9151] ]      # F


def rescale(T0,T1):
    return (T1/T0)**3 * np.exp(6400*(1/T0-1/T1))


def getDarkStats(darke, spike, temp, dc_warm):
    dark = fits.getdata(darke).astype(float)
    header = fits.getheader(darke)
    systemp = float(header['SYSTEMP'])
    temp0 = systemp

    print('\n\n')
    print('> darke:   {}'.format(darke))
    print('> spike:   {} e-/s/pix'.format(spike))
    print('> systemp: {} K'.format(temp0))
    print('> warm pixel dc: {} e-/s/pix'.format(dc_warm))
    
    dc_pk = spike
    if temp > 0:
        T0 = temp0
        T1 = temp
        print('> temp:    {} K'.format(temp))
        dark = dark*rescale(T0,T1)

    idx_warm = np.logical_and(dark>dc_pk, dark<=25)   # get the total number of warm pixels
    idx_hot  = dark > 25
    idx0 = np.logical_and(dark>dc_pk, dark<=25)
    idx1 = np.logical_and(dark>dc_pk, dark<=dc_warm)
    print('--> hot pixel number: {}'.format(np.sum(idx_hot)))
    print('--> warm pixel number [> {}e-/s/pix & < 25e-/s/pix]: {}'.format(dc_warm, np.sum(idx_warm)))
    

    frac0 = np.sum(idx0)/dark.shape[0]/dark.shape[1]
    frac1 = np.sum(idx1)/dark.shape[0]/dark.shape[1]
    print('==> [warm pixels not excluded]')
    print('\tmean dc: {:.4f} e-/s/pix'.format(np.mean(dark[dark<25])))
    print('\tfrac[>{:.4f} e-/s/pix] = {:.3e}'.format(dc_pk, frac0))
    print('==> [warm pixels excluded]')
    print('\tmean dc: {:.4f} e-/s/pix'.format(np.mean(dark[dark<dc_warm])))
    print('\tfrac[>{:.4f} e-/s/pix] = {:.3e}'.format(dc_pk, frac1))


def getDarkStats_for_regions(darke, spike, temp, dc_warm):
    dark = fits.getdata(darke).astype(float)
    header = fits.getheader(darke)
    systemp = float(header['SYSTEMP'])
    temp0 = systemp

    print('> darke:   {}'.format(darke))
    print('> spike:   {} e-/s/pix'.format(spike))
    print('> systemp: {} K'.format(temp0))
    print('> warm pixel dc: {} e-/s/pix'.format(dc_warm))
    
    dc_pk = spike
    if temp > 0:
        T0 = temp0
        T1 = temp
        print('> temp:    {} K'.format(temp))
        dark = dark*rescale(T0,T1)

    # loop over regions
    for i in range(len(reg)):
        print('---------------- {} ----------------'.format(reg_names[i]))
        # get region coordiantes
        py0, py1, px0, px1 = reg[i]

        # get darke in that region
        regdarke = dark[px0:px1, py0:py1]
        print('--> total number of pixels: <<< {} >>>'.format(regdarke.shape[0]*regdarke.shape[1]))

        idx_warm = np.logical_and(regdarke>dc_pk, regdarke<=25)   # get the total number of warm pixels
        idx_hot  = dark > 25
        idx0 = np.logical_and(regdarke>dc_pk, regdarke<=25)
        idx1 = np.logical_and(regdarke>dc_pk, regdarke<=dc_warm)
        print('--> hot pixel number: {}'.format(np.sum(idx_hot)))
        print('--> warm pixel number [> {}e-/s/pix & < 25e-/s/pix]: {}'.format(dc_warm, np.sum(idx_warm)))
        

        frac0 = np.sum(idx0)/regdarke.shape[0]/regdarke.shape[1]
        frac1 = np.sum(idx1)/regdarke.shape[0]/regdarke.shape[1]
        print('==> [warm pixels not excluded]')
        print('\tmean dc: {:.4f} e-/s/pix'.format(np.mean(regdarke[regdarke<25])))
        print('\tfrac[>{:.4f} e-/s/pix] = {:.3e}'.format(dc_pk, frac0))
        print('==> [warm pixels excluded]')
        print('\tmean dc: {:.4f} e-/s/pix'.format(np.mean(regdarke[regdarke<dc_warm])))
        print('\tfrac[>{:.4f} e-/s/pix] = {:.3e}'.format(dc_pk, frac1))

def main():
    parser = ArgumentParser()
    parser.add_argument('--darke', '-d', default=None, type=str,
                        help='filename of processed darke.fits')
    parser.add_argument('--spike', '-p', default=0.08, type=float,
                        help='dark spike threshold value (e-/s/pix)')
    parser.add_argument('--temp', '-t', default=-1, type=float,
                        help='target temperature of darke to be rescaled to')
    parser.add_argument('--dc_warm', '-w', default=1, type=float,
                        help='minimum dark current of warm pixel')
    parser.add_argument('--regions', '-r', action='store_true',
                        help='If true, then print dark current statistics for each region')
    args = parser.parse_args()

    if args.regions is False:
        getDarkStats(args.darke, args.spike, args.temp, args.dc_warm)
    else:
        getDarkStats_for_regions(args.darke, args.spike, args.temp, args.dc_warm)
        getDarkStats(args.darke, args.spike, args.temp, args.dc_warm)


if __name__=='__main__':
    main()
